package tsj.ai.alibaba.data.analysis.service.vectorstore;

import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
import tsj.ai.alibaba.data.analysis.common.connector.accessor.Accessor;
import tsj.ai.alibaba.data.analysis.common.connector.accessor.AccessorFactory;
import tsj.ai.alibaba.data.analysis.common.connector.bo.DbQueryParameter;
import tsj.ai.alibaba.data.analysis.common.connector.bo.ForeignKeyInfoBO;
import tsj.ai.alibaba.data.analysis.common.connector.config.DbConfig;
import tsj.ai.alibaba.data.analysis.common.request.AgentSearchRequest;
import tsj.ai.alibaba.data.analysis.common.request.SchemaInitRequest;
import tsj.ai.alibaba.data.analysis.constant.Constant;
import tsj.ai.alibaba.data.analysis.service.TableMetadataService;
import tsj.ai.alibaba.data.analysis.common.connector.bo.TableInfoBO;

import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.stream.Collectors;

import static tsj.ai.alibaba.data.analysis.util.DocumentConverterUtil.convertColumnsToDocuments;
import static tsj.ai.alibaba.data.analysis.util.DocumentConverterUtil.convertTablesToDocuments;

/**
 * @author taoshujian
 * @version 1.0
 * @since 2025-11-13 17:41:09
 */
@Slf4j
@Service
public class AgentVectorStoreServiceImpl  implements AgentVectorStoreService {

    private static final String DEFAULT = "default";

    private final VectorStore vectorStore; // 由Spring Boot自动配置

    private final ExecutorService dbOperationExecutor;

    /**
     * 相似度阈值配置，用于过滤相似度分数大于等于此阈值的文档
     */
    @Value("${spring.ai.vectorstore.similarityThreshold:0.2}")
    protected double similarityThreshold;

    @Value("${spring.ai.vectorstore.batch-topk-limit:5000}")
    protected int batchTopkLimit;

    @Value("${spring.ai.vectorstore.agent-query-topk-limit:25}")
    protected int agentQueryTopkLimit;

    protected final AccessorFactory accessorFactory;

    protected final BatchingStrategy batchingStrategy;

    protected final TableMetadataService tableMetadataService;

    public AgentVectorStoreServiceImpl(VectorStore vectorStore, ExecutorService dbOperationExecutor,
                                       BatchingStrategy batchingStrategy, AccessorFactory accessorFactory,
                                       TableMetadataService tableMetadataService) {
        this.vectorStore = vectorStore;
        this.dbOperationExecutor = dbOperationExecutor;
        this.batchingStrategy = batchingStrategy;
        this.accessorFactory = accessorFactory;
        this.tableMetadataService = tableMetadataService;
        log.info("VectorStore type: {}", vectorStore.getClass().getSimpleName());
    }

    /**
     * 查询所有向量数据
     * @return 所有文档列表
     */
    @Override
    public List<Document> findAllDocuments() {
        try {
            // 构建一个无过滤条件的搜索请求，获取所有文档
            org.springframework.ai.vectorstore.SearchRequest searchRequest =
                    org.springframework.ai.vectorstore.SearchRequest.builder()
                            //.query(DEFAULT) // 使用默认查询字符串
                            .topK(Integer.MAX_VALUE) // 设置最大返回数量
                            .similarityThreshold(0.0) // 设置最低相似度阈值以获取所有文档
                            .build();

            List<Document> results = vectorStore.similaritySearch(searchRequest);
            log.info("一共有文档 {}个", results.size());
            return results;
        } catch (Exception e) {
            log.error("Error while fetching all documents", e);
            throw new RuntimeException("Failed to fetch all documents", e);
        }
    }

    @Override
    public List<Document> search(AgentSearchRequest searchRequest) {
        Assert.notNull(searchRequest, "SearchRequest cannot be null");
        Assert.notNull(searchRequest.getAgentId(), "agentId  cannot be null");
        org.springframework.ai.vectorstore.SearchRequest.Builder builder = org.springframework.ai.vectorstore.SearchRequest
                .builder();

        if (StringUtils.hasText(searchRequest.getQuery()))
            builder.query(searchRequest.getQuery());

        if (Objects.nonNull(searchRequest.getTopK()))
            builder.topK(searchRequest.getTopK());

        String filterFormatted = buildFilterExpressionString(searchRequest.getMetadataFilter());
        log.info("向量搜索的查询条件: {}", filterFormatted);
        if (StringUtils.hasText(filterFormatted))
            builder.filterExpression(filterFormatted);
        builder.similarityThreshold(similarityThreshold);
        List<Document> results = vectorStore.similaritySearch(builder.build());
        log.info("Search completed. Found {} documents for SearchRequest: {}", results.size(), searchRequest);
        return results;

    }

    @Override
    public Boolean schema(String agentId, SchemaInitRequest schemaInitRequest) throws Exception {
        log.info("Starting schema initialization for agent: {}", agentId);
        DbConfig config = schemaInitRequest.getDbConfig();
        DbQueryParameter dqp = DbQueryParameter.from(config)
                .setSchema(config.getSchema())
                .setTables(schemaInitRequest.getTables());

        try {
            // 根据当前DbConfig获取Accessor
            Accessor dbAccessor = accessorFactory.getAccessorByDbConfig(config);

            // 清理旧数据
            log.info("Clearing existing schema data for agent: {}", agentId);
            clearSchemaDataForAgent(agentId);
            log.debug("Successfully cleared existing schema data for agent: {}", agentId);

            // 处理外键
            log.debug("Fetching foreign keys for agent: {}", agentId);
            List<ForeignKeyInfoBO> foreignKeys = dbAccessor.showForeignKeys(config, dqp);
            log.info("Found {} foreign keys for agent: {}", foreignKeys.size(), agentId);

            Map<String, List<String>> foreignKeyMap = buildForeignKeyMap(foreignKeys);
            log.debug("Built foreign key map with {} entries for agent: {}", foreignKeyMap.size(), agentId);

            // 处理表和列
            log.debug("Fetching tables for agent: {}", agentId);
            List<TableInfoBO> tables = dbAccessor.fetchTables(config, dqp);
            log.info("Found {} tables for agent: {}", tables.size(), agentId);

            if (tables.size() > 5) {
                // 对于大量表，使用并行处理
                log.info("Processing {} tables in parallel mode for agent: {}", tables.size(), agentId);
                tableMetadataService.batchEnrichTableMetadata(tables, config, foreignKeyMap);
            }
            else {
                // 对于少量表，使用批量处理
                log.info("Processing {} tables in batch mode for agent: {}", tables.size(), agentId);
                tableMetadataService.batchEnrichTableMetadata(tables, config, foreignKeyMap);
            }

            log.info("Successfully processed all tables for agent: {}", agentId);

            // 转换为文档
            List<Document> columnDocs = convertColumnsToDocuments(agentId, tables);
            List<Document> tableDocs = convertTablesToDocuments(agentId, tables);

            // 存储文档
            log.info("Storing {} columns and {} tables for agent: {}", columnDocs.size(), tableDocs.size(), agentId);
            storeSchemaDocuments(columnDocs, tableDocs);
            log.info("Successfully stored all documents for agent: {}", agentId);

            // 查询所有向量数据库

            return true;
        }
        catch (Exception e) {
            log.error("Failed to process schema for agent: {}", agentId, e);
            return false;
        }
    }

    protected Map<String, List<String>> buildForeignKeyMap(List<ForeignKeyInfoBO> foreignKeys) {
        Map<String, List<String>> map = new HashMap<>();
        for (ForeignKeyInfoBO fk : foreignKeys) {
            String key = fk.getTable() + "." + fk.getColumn() + "=" + fk.getReferencedTable() + "."
                    + fk.getReferencedColumn();

            map.computeIfAbsent(fk.getTable(), k -> new ArrayList<>()).add(key);
            map.computeIfAbsent(fk.getReferencedTable(), k -> new ArrayList<>()).add(key);
        }
        return map;
    }

    @Override
    public Boolean deleteDocumentsByVectorType(String agentId, String vectorType) throws Exception {
        return null;
    }

    @Override
    public Boolean deleteDocumentsByMetedata(String agentId, Map<String, Object> metadata) throws Exception {
        return null;
    }

    /**
     * 并行处理表元数据，提高大量表时的处理性能
     * @param tables 表列表
     * @param config 数据库配置
     * @param foreignKeyMap 外键映射
     * @throws Exception 处理失败时抛出异常
     */
    private void processTablesInParallel(List<TableInfoBO> tables, DbConfig config,
                                         Map<String, List<String>> foreignKeyMap) throws Exception {


    }

    /**
     * 将列表分成指定大小的子列表
     * @param list 原始列表
     * @param batchSize 批次大小
     * @param <T> 列表元素类型
     * @return 分批后的列表
     */
    private <T> List<List<T>> partitionList(List<T> list, int batchSize) {
        List<List<T>> partitions = new ArrayList<>();
        for (int i = 0; i < list.size(); i += batchSize) {
            partitions.add(list.subList(i, Math.min(i + batchSize, list.size())));
        }
        return partitions;
    }

    protected void storeSchemaDocuments(List<Document> columns, List<Document> tables) {
        // 串行去批写入，并行流的时候有API限速了
        List<List<Document>> columnBatches = batchingStrategy.batch(columns);
        columnBatches.forEach(vectorStore::add);

        List<List<Document>> tableBatches = batchingStrategy.batch(tables);
        tableBatches.forEach(vectorStore::add);

    }


    protected void clearSchemaDataForAgent(String agentId) throws Exception {
        deleteDocumentsByVectorType(agentId, Constant.COLUMN);
        deleteDocumentsByVectorType(agentId, Constant.TABLE);
    }


    private void batchDelDocumentsWithFilter(String filterExpression) {
        Set<String> seenDocumentIds = new HashSet<>();
        // 分批获取，因为Milvus等向量数据库的topK有限制
        List<Document> batch;
        int newDocumentsCount;
        int totalDeleted = 0;

        do {
            batch = vectorStore.similaritySearch(org.springframework.ai.vectorstore.SearchRequest.builder()
                    .query(DEFAULT)// 使用默认的查询字符串，因为有的嵌入模型不支持空字符串
                    .filterExpression(filterExpression)
                    .similarityThreshold(0.0)// 设置最低相似度阈值以获取元数据匹配的所有文档
                    .topK(batchTopkLimit)
                    .build());

            // 过滤掉已经处理过的文档，只删除未处理的文档
            List<String> idsToDelete = new ArrayList<>();
            newDocumentsCount = 0;

            for (Document doc : batch) {
                if (seenDocumentIds.add(doc.getId())) {
                    // 如果add返回true，表示这是一个新的文档ID
                    idsToDelete.add(doc.getId());
                    newDocumentsCount++;
                }
            }

            // 删除这批新文档
            if (!idsToDelete.isEmpty()) {
                vectorStore.delete(idsToDelete);
                totalDeleted += idsToDelete.size();
            }

        }
        while (newDocumentsCount > 0); // 只有当获取到新文档时才继续循环

        log.info("Deleted {} documents with filter expression: {}", totalDeleted, filterExpression);
    }

    /**
     * 构建过滤表达式字符串，目前FilterExpressionBuilder 不支持链式拼接元数据过滤，所以只能使用字符串拼接
     * @param filterMap
     * @return
     */
    protected final String buildFilterExpressionString(Map<String, Object> filterMap) {
        if (filterMap == null || filterMap.isEmpty()) {
            return null;
        }

        // 验证键名是否合法（只包含字母、数字和下划线）
        for (String key : filterMap.keySet()) {
            if (!key.matches("[a-zA-Z_][a-zA-Z0-9_]*")) {
                throw new IllegalArgumentException("Invalid key name: " + key
                        + ". Keys must start with a letter or underscore and contain only alphanumeric characters and underscores.");
            }
        }

        return filterMap.entrySet().stream().map(entry -> {
            String key = entry.getKey();
            Object value = entry.getValue();

            // 处理空值
            if (value == null) {
                return key + " == null";
            }

            // 根据值的类型决定如何格式化
            if (value instanceof String) {
                // 转义字符串中的特殊字符
                String escapedValue = escapeStringLiteral((String) value);
                return key + " == '" + escapedValue + "'";
            }
            else if (value instanceof Number) {
                // 数字类型直接使用
                return key + " == " + value;
            }
            else if (value instanceof Boolean) {
                // 布尔值使用小写形式
                return key + " == " + ((Boolean) value).toString().toLowerCase();
            }
            else if (value instanceof Enum) {
                // 枚举类型，转换为字符串并转义
                String enumValue = ((Enum<?>) value).name();
                String escapedValue = escapeStringLiteral(enumValue);
                return key + " == '" + escapedValue + "'";
            }
            else {
                // 其他类型尝试转换为字符串并转义
                String stringValue = value.toString();
                String escapedValue = escapeStringLiteral(stringValue);
                return key + " == '" + escapedValue + "'";
            }
        }).collect(Collectors.joining(" && "));
    }

    /**
     * 转义字符串字面量中的特殊字符
     */
    private String escapeStringLiteral(String input) {
        if (input == null) {
            return "";
        }

        // 转义反斜杠和单引号
        String escaped = input.replace("\\", "\\\\").replace("'", "\\'");

        // 转义其他特殊字符
        escaped = escaped.replace("\n", "\\n")
                .replace("\r", "\\r")
                .replace("\t", "\\t")
                .replace("\b", "\\b")
                .replace("\f", "\\f");

        return escaped;
    }

    @Override
    public List<Document> getDocumentsForAgent(String agentId, String query, String vectorType) {
        AgentSearchRequest searchRequest = AgentSearchRequest.getInstance(agentId);
        searchRequest.setQuery(query);
        searchRequest.setTopK(agentQueryTopkLimit);
        searchRequest.setMetadataFilter(Map.of(Constant.VECTOR_TYPE, vectorType));

        return search(searchRequest);
    }

    @Override
    public boolean hasDocuments(String agentId) {

        return false;
    }

    @Override
    public void addDocuments(String agentId, List<Document> documents) {

    }

    @Override
    public int estimateDocuments(String agentId) {
        return 0;
    }

    /**
     * 销毁方法，在Spring容器销毁bean时调用，用于关闭线程池
     */
    @PreDestroy
    public void destroy() {
        if (dbOperationExecutor != null && !dbOperationExecutor.isShutdown()) {
            log.info("Shutting down database operation executor");
            dbOperationExecutor.shutdown();
            try {
                // 等待正在执行的任务完成，最多等待30秒
                if (!dbOperationExecutor.awaitTermination(30, java.util.concurrent.TimeUnit.SECONDS)) {
                    log.warn("Database operation executor did not terminate gracefully, forcing shutdown");
                    dbOperationExecutor.shutdownNow();
                }
            }
            catch (InterruptedException e) {
                log.warn("Interrupted while waiting for database operation executor to terminate");
                dbOperationExecutor.shutdownNow();
                Thread.currentThread().interrupt();
            }
        }
    }




}
